
import glob
import multiprocessing
import os

import mindspore as ms
from PIL import Image
from mindspore import nn
import mindspore.dataset as ds
from day2_ChangeFormerV6.dataset.CD_dataset import CDDataset
from day2_ChangeFormerV6.models.ChangeFormer import ChangeFormerV6
from day2_ChangeFormerV6.utils.custom_with_cell import CustomWithEvalCell
from day2_ChangeFormerV6.utils.loss import CrossEntropyWithLogits
from mindspore.common.initializer import One, Normal
from mindspore import context
from day2_ChangeFormerV6.dataset.cd_dataset_new import create_dataset_new
import numpy as np
from mindspore import ops
from mindspore import save_checkpoint

ms.set_context(mode=ms.GRAPH_MODE,device_target="CPU", pynative_synchronize=True)

from mindspore import dtype as mstype

def create_dataset(batch = 16):
    # instanceof dataset
    dataset_generator = CDDataset(root_dir=r'E:\ChangeFormer模型迁移相关内容\迁移准备\LEVIR-CD-256', img_size=256, split='val', is_train=True, label_transform='norm')

    dataset = ds.GeneratorDataset(dataset_generator, ["img_A", "img_B", "label"], shuffle=False)
    dataset = dataset.batch(batch_size=batch)
    return dataset

def get_data(num):
    """生成样本数据及对应的标签"""
    for _ in range(num):
        imgA = ms.Tensor(shape=(1, 3, 256, 256), dtype=mstype.float32, init=Normal())
        imgB = ms.Tensor(shape=(1, 3, 256, 256), dtype=mstype.float32, init=Normal())
        label = ms.Tensor(shape=(1, 1, 256, 256), dtype=mstype.uint32, init=Normal())
        yield imgA, imgB, label


def get_imgset(root, mode, batch_size=2):
    files_A = sorted(glob.glob(os.path.join(root, '%s/A' % mode) + '/*.*'))
    files_B = sorted(glob.glob(os.path.join(root, '%s/B' % mode) + '/*.*'))
    files_Label = sorted(glob.glob(os.path.join(root, '%s/label' % mode) + '/*.*'))


    for img_index in range(0, len(files_A), batch_size):

        for i in range(batch_size):
            input_A = Image.open(files_A[(img_index+i) % len(files_A)]).convert('RGB')
            input_B = Image.open(files_B[(img_index+i) % len(files_B)]).convert('RGB')
            label = Image.open(files_Label[(img_index+i) % len(files_Label)]).convert('L')

            imgA = np.asarray(input_A).astype(np.float32)
            imgB = np.asarray(input_B).astype(np.float32)
            label = np.asarray(label).astype(np.uint8)

            # imgA = ops.expand_dims(ms.Tensor(imgA / 255.).transpose(-1, 0, 1), 0)
            imgA = ms.Tensor(imgA / 255.).transpose(-1, 0, 1)
            # imgB = ops.expand_dims(ms.Tensor(imgB / 255.).transpose(-1, 0, 1), 0)
            imgB = ms.Tensor(imgB / 255.).transpose(-1, 0, 1)
            # label = ops.expand_dims(ms.Tensor(label // 255), 0)
            label = ms.Tensor(label // 255)

            # yield imgA_batch, imgB_batch, label_batch
            yield imgA, imgB, label


net = ChangeFormerV6(embed_dim=256)
loss_func = CrossEntropyWithLogits()
custom_eval_net = CustomWithEvalCell(net, loss_func)
custom_eval_net.set_train(False)


# 评价指标
loss = nn.Loss()
acc = nn.Accuracy()
f1_score = nn.F1()
rec = nn.Recall()

# 指标初始化
loss.clear()
acc.clear()
f1_score.clear()
rec.clear()


# 测试使用
# train_dataset = ds.GeneratorDataset(list(get_data(2)), column_names=['img_A', 'img_B', 'label'])
# val_dataset = ds.GeneratorDataset(list(get_imgset(root=r'D:\git_test\Levir_sample', mode='val', batch_size=1)),
#                                     column_names=['imgA', 'imgB', 'label'])
# train_dataset = val_dataset.batch(1)
train_dataset = create_dataset_new(batch=1, mode='val')

expand_dim = ops.ExpandDims()
step = 1
for d in train_dataset.create_dict_iterator():
    label = d['label']
    outputs = custom_eval_net(d['imgA'], d['imgB'], d['label'])
    pred_loss = ms.Tensor(outputs[0])
    loss.update(pred_loss)
    pred = outputs[-1]
    acc.update(pred, label)
    f1_score.update(pred, label)
    rec.update(pred, label)
    print('step:', step)
    step += 1

# 评估结果
mLoss = loss.eval()
mAcc = acc.eval()
mF1_score = f1_score.eval()
mRec = rec.eval()
print('mLoss: ', mLoss)
print('mAcc: ', mAcc)
print('mRec: ', mRec)
print('mF1-score: ', mF1_score)



